import type { Sequential, Tensor, Logs, History } from "@tensorflow/tfjs-node";

async function trainModel(
  model: Sequential,
  trainingFeatureTensor: Tensor,
  trainingLabelTensor: Tensor
): Promise<History> {
  return model.fit(trainingFeatureTensor, trainingLabelTensor, {
    batchSize: 32,
    epochs: 20,
    validationSplit: 0.2,
    callbacks: {
      onEpochEnd: (epoch: number, log: Logs | undefined) => {
        console.log(`Epoch ${epoch} with loss: ${(log as Logs).loss}`);
      },
    },
  });
}

export { trainModel };
